// Copyright [2022] <Tianyu Zhu, Shizhen Zhao, Zihang Ren>

#include <algorithm>
#include <cstdlib>  // Header file needed to use srand and rand
#include <iostream>
#include <stack>
#include <unordered_map>

#include "double_pulse.h"

DoublePulse::~DoublePulse() {
    if (dst_cost_map_) {
        delete[] dst_cost_map_;
        dst_cost_map_ = nullptr;
    }
    if (dst_delay_map_) {
        delete[] dst_delay_map_;
        dst_delay_map_ = nullptr;
    }
    if (src_cost_map_) {
        delete[] src_cost_map_;
        src_cost_map_ = nullptr;
    }
    if (src_delay_map_) {
        delete[] src_delay_map_;
        src_delay_map_ = nullptr;
    }
    if (src_visited_) {
        delete[] src_visited_;
        src_visited_ = nullptr;
    }
    if (dst_visited_) {
        delete[] dst_visited_;
        dst_visited_ = nullptr;
    }
    if (info_) {
        delete[] info_;
        info_ = nullptr;
    }
}

void DoublePulse::SetupTopology(Graph* graph) {
    graph_ = graph;
    num_nodes_ = graph_->NodeSize();
    dst_cost_map_ = new int[num_nodes_];
    dst_delay_map_ = new int[num_nodes_];
    src_cost_map_ = new int[num_nodes_];
    src_delay_map_ = new int[num_nodes_];
    src_visited_ = new bool[num_nodes_];
    dst_visited_ = new bool[num_nodes_];
    info_ = new PulseInfo[num_nodes_]();
}

Path DoublePulse::FindPath(const Flow& flow) {
    // Initialization
    // do dijkstra from end to every other nodes,
    // from two perspectives--cost and delay
    // store the result of dijkstras in cost_map and delay_map
    flow_ = flow;
    clock_t start_time = clock();
    DstDijkstra(flow_.to, false, dst_cost_map_);
    DstDijkstra(flow_.to, true, dst_delay_map_);
    SrcDijkstra(flow_.from, false, src_cost_map_);
    SrcDijkstra(flow_.from, true, src_delay_map_);
    clock_t end_time = clock();
    std::cout << "Pulse initialization time: "
              << double(end_time - start_time) / CLOCKS_PER_SEC * 1000
              << "(ms).\n";

    // Reset Pulse Algorithm
    int cnt = 0;
    for (int i = 0; i < num_nodes_; ++i) {
        src_visited_[i] = false;
        dst_visited_[i] = false;
        info_[i].Clear();
    }
    total_cost_ = kMaxValue;
    results_.clear();

    // Timer starts
    start_time = clock();
    // Start pulse algorithm
    Dfs();
    // Print time elapse
    end_time = clock();
    std::cout << "Pulse time: "
              << double(end_time - start_time) / CLOCKS_PER_SEC * 1000
              << "(ms).\n";
    return results_;
}

// dijkstra algorithm
// if(flag == false)--search according to cost
// if(flag == true)--search according to delay
void DoublePulse::DstDijkstra(
    NodeId dst_node, bool flag, int* node_to_dst_dis) {
    for (int i = 0; i < num_nodes_; ++i) {
        node_to_dst_dis[i] = kMaxValue;
    }
    node_to_dst_dis[dst_node] = 0;
    std::priority_queue<Node<NodeId>, std::vector<Node<NodeId>>> heap;
    std::vector<bool> visited(num_nodes_, false);
    heap.push(Node<NodeId>(dst_node, 0));
    while (!heap.empty()) {
        NodeId u = heap.top().id;
        int weight_u_to_dst = heap.top().distance;
        heap.pop();
        if (visited[u]) continue;
        visited[u] = true;
        for (Link* link : graph_->GetIngressLinks(u)) {
            NodeId v = link->source_id;
            if (visited[v]) continue;
            int weight = flag ? link->delay : link->cost;
            if (node_to_dst_dis[v] > weight_u_to_dst + weight) {
                node_to_dst_dis[v] = weight_u_to_dst + weight;
                heap.push(Node<NodeId>(v, node_to_dst_dis[v]));
            }
        }
    }
}

void DoublePulse::SrcDijkstra(
    NodeId src_node, bool flag, int* node_to_src_dis) {
    for (int i = 0; i < num_nodes_; ++i) {
        node_to_src_dis[i] = kMaxValue;
    }
    node_to_src_dis[src_node] = 0;
    std::priority_queue<Node<NodeId>, std::vector<Node<NodeId>>> heap;
    std::vector<bool> visited(num_nodes_, false);
    heap.push(Node<NodeId>(src_node, 0));
    while (!heap.empty()) {
        NodeId u = heap.top().id;
        int weight_u_to_src = heap.top().distance;
        heap.pop();
        if (visited[u]) continue;
        visited[u] = true;
        for (Link* link : graph_->GetEgressLinks(u)) {
            NodeId v = link->ter_id;
            if (visited[v]) continue;
            int weight = flag ? link->delay : link->cost;
            if (node_to_src_dis[v] > weight_u_to_src + weight) {
                node_to_src_dis[v] = weight_u_to_src + weight;
                heap.push(Node<NodeId>(v, node_to_src_dis[v]));
            }
        }
    }
}

// update the path node U records
// according to the rule mentioned in Chapter 4 of 2012.10 paper
bool DoublePulse::updateSrcPath(NodeId u, int cost, int delay) {
    auto it = info_[u].src_delay_to_cost_map.find(delay);
    if (it == info_[u].src_delay_to_cost_map.end()) {
        info_[u].src_delay_to_cost_map[delay] = cost;
        return true;
    }
    if (it->second > cost) {
        info_[u].src_delay_to_cost_map[delay] = cost;
        return true;
    }
    return false;
}

bool DoublePulse::updateDstPath(NodeId u, int cost, int delay) {
    auto it = info_[u].dst_delay_to_cost_map.find(delay);
    if (it == info_[u].dst_delay_to_cost_map.end()) {
        info_[u].dst_delay_to_cost_map[delay] = cost;
        return true;
    }
    if (it->second > cost) {
        info_[u].dst_delay_to_cost_map[delay] = cost;
        return true;
    }
    return false;
}

std::vector<Link*> DoublePulse::GetOrderedEgressLinks(NodeId node, int delay) {
    std::vector<Node<Link*>> links;
    for (Link* link : graph_->GetEgressLinks(node)) {
        NodeId ter = link->ter_id;
        if (src_visited_[ter]) {
            continue;
        }
        int cost = 10000000;
        for (int d = flow_.delay_lb - delay - link->delay;
             d <= flow_.delay_ub - delay - link->delay; ++d) {
            auto it = info_[ter].dst_delay_to_cost_map.find(d);
            if (it != info_[ter].dst_delay_to_cost_map.end() &&
                it->second < cost) {
                cost = it->second < cost;
            }
        }
        cost += link->cost;
        links.emplace_back(link, cost);
    }
    std::sort(links.begin(), links.end());
    std::vector<Link*> sorted_links;
    // std::cout << "-------\n";
    for (const Node<Link*>& x : links) {
        // std::cout << x.distance << std::endl;
        sorted_links.push_back(x.id);
    }
    return sorted_links;
}

std::vector<Link*> DoublePulse::GetOrderedIngressLinks(NodeId node, int delay) {
    std::vector<Node<Link*>> links;
    for (Link* link : graph_->GetIngressLinks(node)) {
        NodeId ter = link->source_id;
        if (dst_visited_[ter]) {
            continue;
        }
        int cost = 10000000;
        for (int d = flow_.delay_lb - delay - link->delay;
             d <= flow_.delay_ub - delay - link->delay; ++d) {
            auto it = info_[ter].src_delay_to_cost_map.find(d);
            if (it != info_[ter].src_delay_to_cost_map.end() &&
                it->second < cost) {
                cost = it->second < cost;
            }
        }
        cost += link->cost;
        links.emplace_back(link, cost);
    }
    std::sort(links.begin(), links.end());
    std::vector<Link*> sorted_links;
    for (const Node<Link*>& x : links) {
        sorted_links.push_back(x.id);
    }
    return sorted_links;
}

// dfs from start to end, use recursion
void DoublePulse::Dfs() {
    struct Snapshot {
        NodeId node;
        int cost;
        int delay;
        int stage;
    };
    std::stack<Snapshot> src_stack;
    std::stack<Snapshot> dst_stack;
    Path src_path;
    src_path.path_info = {};
    src_path.cost = 0;
    src_path.delay = 0;
    Snapshot snapshot;
    snapshot.node = flow_.from;
    snapshot.cost = 0;
    snapshot.delay = 0;
    snapshot.stage = 0;
    src_stack.push(snapshot);
    Path dst_path;
    dst_path.path_info = {};
    dst_path.cost = 0;
    dst_path.delay = 0;
    snapshot.node = flow_.to;
    dst_stack.push(snapshot);
    int num_iterations = 0;
    while (!src_stack.empty() && !dst_stack.empty()) {
        ++num_iterations;
        Snapshot src_snap = src_stack.top();
        src_stack.pop();
        NodeId u = src_snap.node;
        if (src_snap.stage == 1) {
            src_path.path_info.pop_back();
            src_path.cost -= src_snap.cost;
            src_path.delay -= src_snap.delay;
            src_visited_[u] = false;
        } else {
            int path_cost = src_path.cost + src_snap.cost;
            int path_delay = src_path.delay + src_snap.delay;
            if (u == flow_.to) {
                // src_path.Print();
                if (path_delay <= flow_.delay_ub &&
                    path_delay >= flow_.delay_lb &&
                    path_cost < total_cost_) {
                    // total_cost_ = cost * 0.8;
                    total_cost_ = path_cost;
                    results_ = src_path;
                    results_.path_info.push_back(u);
                    results_.cost = path_cost;
                    results_.delay = path_delay;
                }
            } else {
                // prune
                if (dst_cost_map_[u] + path_cost < total_cost_ &&
                    dst_delay_map_[u] + path_delay <= flow_.delay_ub) {
                    updateSrcPath(u, path_cost, path_delay);
                    std::vector<Link*> egress_links =
                        GetOrderedEgressLinks(u, path_delay);
                    if (!egress_links.empty()) {
                        src_path.path_info.push_back(u);
                        src_path.cost = path_cost;
                        src_path.delay = path_delay;
                        src_visited_[u] = true;
                        src_snap.stage = 1;
                        src_stack.push(src_snap);
                        for (Link* link : egress_links) {
                            NodeId node_id = link->ter_id;
                            Snapshot snap;
                            snap.node = node_id;
                            snap.delay = link->delay;
                            snap.cost = link->cost;
                            snap.stage = 0;
                            src_stack.push(snap);
                        }
                    }
                }
            }
        }

        Snapshot dst_snap = dst_stack.top();
        dst_stack.pop();
        NodeId v = dst_snap.node;
        if (dst_snap.stage == 1) {
            dst_path.path_info.pop_back();
            dst_path.cost -= dst_snap.cost;
            dst_path.delay -= dst_snap.delay;
            dst_visited_[v] = false;
        } else {
            int path_cost = dst_path.cost + dst_snap.cost;
            int path_delay = dst_path.delay + dst_snap.delay;
            if (v == flow_.from) {
                // dst_path.Print();
                if (path_delay <= flow_.delay_ub &&
                    path_delay >= flow_.delay_lb &&
                    path_cost < total_cost_) {
                    // total_cost_ = cost * 0.8;
                    total_cost_ = path_cost;
                    results_ = dst_path;
                    results_.path_info.push_back(v);
                    results_.cost = path_cost;
                    results_.delay = path_delay;
                }
            } else {
                // prune
                if (src_cost_map_[v] + path_cost < total_cost_ &&
                    src_delay_map_[v] + path_delay <= flow_.delay_ub) {
                    updateDstPath(v, path_cost, path_delay);
                    std::vector<Link*> ingress_links =
                        GetOrderedIngressLinks(v, path_delay);
                    if (!ingress_links.empty()) {
                        dst_path.path_info.push_back(v);
                        dst_path.cost = path_cost;
                        dst_path.delay = path_delay;
                        dst_visited_[v] = true;
                        dst_snap.stage = 1;
                        dst_stack.push(dst_snap);
                        for (Link* link : ingress_links) {
                            NodeId node_id = link->source_id;
                            Snapshot snap;
                            snap.node = node_id;
                            snap.delay = link->delay;
                            snap.cost = link->cost;
                            snap.stage = 0;
                            dst_stack.push(snap);
                        }
                    }
                }
            }
        }
    }
    std::cout << "Total number of iterations: " << num_iterations << "\n";
}

void DoublePulse::print() const {
    results_.Print();
}
